package cn.edu.dutir.searchengine;

import java.io.BufferedReader;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.PrintStream;
import java.io.Reader;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * Retrieval performance evaluator.
 * 
 */
public class Evaluator {

	/**
	 * Relevant(standard or reference) set.
	 */
	Set<String> mRelevenceSet = null;
	/**
	 * Answer set,that is the result set returned from a retrieval system.
	 */
	List<String> mAnswerList = null;

	/**
	 * Store the precision value at different rank.
	 */
	List<Double> mPrecisionAtRank = null;
	/**
	 * Store the recall value at different rank.
	 */
	List<Double> mRecallAtRank = null;
	/**
	 * Store the 11-point precision-recall value. Here we only store the
	 * precision values because the recall values can easily be obtained using
	 * the array index.
	 */
	List<Double> mTenPointPRCure = null;

	public static File mRelevantFile = null;
	public static File mAnswerFile = null;
	/**
	 * Construct a retrieval performance evaluator.
	 */
	public Evaluator() {
		mRelevenceSet = new HashSet<String>();
		mAnswerList = new ArrayList<String>();
		mPrecisionAtRank = new ArrayList<Double>();
		mRecallAtRank = new ArrayList<Double>();
		mTenPointPRCure = new ArrayList<Double>(10);
	}

	/**
	 * Evaluate the system's retrieval performance with regard to to a query,
	 * using the precision-at-recall values.
	 * <p>
	 * <b>Note</b>: mean average precision are query-based.
	 * 
	 * @param inRelevanceFile The file that stores the relevant set.
	 * @param inAnswerFile The file that stores the answer set.
	 * @throws IOException
	 */
	public void evaluate(File inRelevanceFile, File inAnswerFile)
			throws IOException {
		loadCollection(mRelevenceSet, inRelevanceFile);
		loadCollection(mAnswerList, inAnswerFile);

		int curRelevanceDocNum = 0;
		int totalRelevanceDoctNum = mRelevenceSet.size();
		int baseRelevanceDocNum = (int)(0.1f * totalRelevanceDoctNum);
		boolean next = false;
		for (int i = 0; i < mAnswerList.size(); i++) {
			if (curRelevanceDocNum >= totalRelevanceDoctNum) {
				break;
			}
			if (mRelevenceSet.contains(mAnswerList.get(i))) {
				curRelevanceDocNum++;
				next = true;
			}
			double precision = curRelevanceDocNum / (i + 1.0);
			double recall = curRelevanceDocNum / (totalRelevanceDoctNum + 0.0);
			mPrecisionAtRank.add(precision);
			mRecallAtRank.add(recall);
			if (next && (curRelevanceDocNum % baseRelevanceDocNum) == 0) {
				mTenPointPRCure.add(precision);
				next = false;
			}
		}
		/**
		 * mean average precision
		 */
		double map = 0.0;
		if (mTenPointPRCure.size() != 0) {
			map = sum(mTenPointPRCure) / mTenPointPRCure.size();
		}
		System.out.println(toString(mTenPointPRCure));
		System.out.printf("MAP: %5.4f \r\n", map);
	}

	/**
	 * Calculate the sum of a double array.
	 */
	public double sum(List<Double> dl) {
		double sum = 0.0;
		for (double v : dl) {
			sum += v;
		}
		return sum;
	}

	/**
	 * Return a formated, string-based representation of the specified double
	 * list.
	 */
	public String toString(List<Double> dl) {
		double base = 0.1;
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		PrintStream ps = new PrintStream(baos);
		ps.append(dl.size() + "-Point Precision-Recall\r\n");
		ps.append(" -------------------- \r\n");
		ps.printf("| %s | %s |\r\n", "Recall", "Precision");
		for (int i = 0; i < dl.size(); i++) {
			ps.printf("| %6.4f | %9.4f |\r\n", base, dl.get(i));
			base += 0.1;
		}
		ps.append(" -------------------- \r\n");
		return new String(baos.toByteArray());
	}

	/**
	 * Reset the evaluator for a new query.
	 */
	public void clear() {
		mRelevenceSet.clear();
		mAnswerList.clear();
		mPrecisionAtRank.clear();
		mRecallAtRank.clear();
	}

	/**
	 * Load a collection of string-based objects from the specified file with
	 * default character encoding. In this implementation, we adopted "utf-8" as
	 * default
	 */
	public static void loadCollection(Collection<String> c, File inFile)
			throws IOException {
		loadCollection(c, inFile, "utf-8");
	}

	/**
	 * Load a collection of string-based objects from the specified file with
	 * the specified character encoding.
	 */
	public static void loadCollection(Collection<String> c, File inFile,
			String charset) throws IOException {
		if (inFile == null || !inFile.exists() || !inFile.isFile()) {
			String msg = "inFile must exist and be a plain file, but now found inFile = "
					+ inFile;
			throw new IllegalArgumentException(msg);
		}
		BufferedReader reader = null;
		InputStreamReader is = null;
		FileInputStream fis = null;
		try {
			fis = new FileInputStream(inFile);
			is = new InputStreamReader(fis, charset);
			reader = new BufferedReader(is);
			loadCollection(c, reader);
		} finally {
			reader.close();
			is.close();
			fis.close();
		}
	}

	/**
	 * Load a collection a string-based objects from the specified stream
	 * <code>reader</code>
	 */
	public static void loadCollection(Collection<String> c, Reader reader)
			throws IOException {
		if (reader == null) {
			String msg = "reader can be null";
			throw new IllegalArgumentException(msg);
		}
		BufferedReader br = null;
		if (!(reader instanceof BufferedReader)) {
			br = (BufferedReader) reader;
		} else {
			br = new BufferedReader(reader);
		}
		String line = "";
		while ((line = br.readLine()) != null) {
			c.add(line.trim());
		}
		c.remove("");
	}

	/**
	 * Parse environment argument passed in from user.
	 * 
	 * @param args Environment argument passed in.
	 */
	private static void parseArgv(String args[]) {
		if (args == null || args.length < 2) {
			usage();
		}
		for (int i = 0; i < args.length; i++) {
			if (args[i].equals("-r")) {
				mRelevantFile = new File(args[++i]);
			} else if (args[i].equals("-a")) {
				mAnswerFile = new File(args[++i]);
			} else {
				usage();
			}
		}
	}
	
	/**
	 * Display usage message.
	 */
	public static void usage() {
		String msg = "Usage:\t Evaluator [-r relevantSet] [-a answerSet]\r\n"
				+ "\t -r relevantSet ---- indicate the path of relevant result file\r\n"
				+ "\t -a answerSet   ---- indicate the path of the answer result file\r\n";
		System.out.println(msg);
		System.exit(1);
	}
	
	/**
	 * The main process
	 * 
	 * @param args
	 * @throws IOException
	 */
	public static void main(String args[]) throws IOException {
		parseArgv(args);
		Evaluator evaluator = new Evaluator();
		evaluator.evaluate(mRelevantFile, mAnswerFile);
	}
}
